# -*- coding: utf-8 -*-

import os
import numpy as np
from PIL import Image
from imgaug import augmenters as iaa
import shutil

"""
 数据增强只有正样本 没有对于负样本的增强，建议解决这个问题
"""

# 路径配置 注意路径末尾必须加/
XML_DIR = "../temp/srcxml/"  # 最终增广文件夹
IMG_DIR = "../temp/srcimg/"

DST_XML_DIR = "../temp/dstxml/"  # 差异化文件夹
DST_IMG_DIR = "../temp/dstimg/"

# 增强后数据目标路径
AUG_XML_DIR = "../output/Annotations/"  # 存储增强后的XML文件夹路径
AUG_IMG_DIR = "../output/JPEGImages/"  # 存储增强后的影像文件夹路径


def data_augmentation():
    # 生成的图片数量
    AUGLOOP = 6
    # 影像增强
    seq = iaa.Sequential([
        iaa.Flipud(0.5),  # vertically flip 20% of all images
        iaa.Fliplr(0.5),  # 镜像
        iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs
        iaa.GaussianBlur(sigma=(0, 0.0)),  # iaa.GaussianBlur(0.5),
        iaa.Affine(
            translate_px={"x": 15, "y": 15},
            scale=(0.8, 0.95),
            rotate=(-30, 30)
        )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
    ])

    for root, sub_folders, files in os.walk(DST_IMG_DIR):
        for name in files:
            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
                # 读取图片
                img = Image.open(os.path.join(DST_IMG_DIR, name[:-4] + '.jpg'))
                img = np.array(img)

                # 存储变化后的图片
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(IMG_DIR, str(name[:-4]) + "_" + str(epoch) + '.jpg')
                # image_auged = bbs.draw_on_image(image_aug, thickness=0)
                # print(str(name[:-4]) + "_aug_" + str(epoch) + '.jpg')
                Image.fromarray(image_aug).save(path)

                # 存储变化后的XML
                xml_path = str(name[:-4]) + "_" + str(epoch) + '.xml'
                new_name = name.split('.')[0]
                # print(xml_path)
                shutil.copyfile(DST_XML_DIR + new_name + ".xml", XML_DIR + xml_path)
    # 将处理后的文件复制到最后的数据目录
    for root, sub_folders, files in os.walk(XML_DIR):
        for name in files:
            # for epoch in range(AUGLOOP):
            new_name = name.split('.')[0]
            # new_name = str(name[:-4]) + "_" + str(epoch) + '.xml'
            print(new_name)
            shutil.copyfile(XML_DIR + new_name + ".xml", AUG_XML_DIR + new_name + ".xml")
            shutil.copyfile(IMG_DIR + new_name + ".jpg", AUG_IMG_DIR + new_name + ".jpg")


if __name__ == '__main__':
    data_augmentation()
